import pymysql
from pymysql import OperationalError

# 配置数据库连接信息
DB_CONFIG = {
    'host': '123.56.171.51',
    'user': 'root',
    'password': 'lgmi6395',
    'db': 'jg_base_new',
    'charset': 'utf8mb4',
    'cursorclass': pymysql.cursors.DictCursor
}

# 需要添加的字段配置（修改这里）
COLUMNS_TO_ADD = [
    {
        'name': 'tenant_id',
        'type': 'varchar',
        'definition': 'varchar(30) COMMENT \'租户ID\''
    },
    {
        'name': 'remark',
        'type': 'varchar',
        'definition': 'varchar(100) COMMENT \'备注\''
    },
    {
        'name': 'deleted_tag',
        'type': 'varchar',
        'definition': 'varchar(30) NOT NULL DEFAULT \'N\' COMMENT \'删除标记\''
    }
]


def get_all_tables(connection):
    """获取数据库中的所有表名"""
    with connection.cursor() as cursor:
        cursor.execute("SHOW TABLES")
        return [row[f'Tables_in_{DB_CONFIG["db"]}'] for row in cursor.fetchall()]


def column_exists(connection, table, column_name):
    """检查字段是否已存在"""
    with connection.cursor() as cursor:
        cursor.execute(f"""
            SELECT COLUMN_NAME 
            FROM INFORMATION_SCHEMA.COLUMNS 
            WHERE TABLE_SCHEMA = %s 
            AND TABLE_NAME = %s 
            AND COLUMN_NAME = %s
        """, (DB_CONFIG['db'], table, column_name))
        return cursor.fetchone() is not None


def add_columns_to_tables():
    try:
        # 建立数据库连接
        connection = pymysql.connect(**DB_CONFIG)

        tables = get_all_tables(connection)
        print(f"找到 {len(tables)} 张表需要处理")

        for table in tables:
            print(f"\n正在处理表: {table}")
            with connection.cursor() as cursor:
                for column in COLUMNS_TO_ADD:
                    # 检查字段是否已存在
                    if column_exists(connection, table, column['name']):
                        print(f"字段 {column['name']} 已存在，跳过")
                        # sql = f"ALTER TABLE {table} MODIFY COLUMN {column['name']} {column['definition']}"
                        continue
                    else:
                        # 执行添加字段操作
                        sql = f"ALTER TABLE `{table}` ADD COLUMN {column['name']} {column['definition']}"
                    print(f"执行SQL: {sql}")
                    cursor.execute(sql)
                connection.commit()

        print("\n所有操作已完成！")

    except OperationalError as e:
        print(f"数据库连接失败: {e}")
    except Exception as e:
        print(f"发生错误: {e}")
        connection.rollback()
    finally:
        if 'connection' in locals() and connection.open:
            connection.close()


if __name__ == "__main__":
    add_columns_to_tables()